# 6 Z字形变换
# https://leetcode.cn/problems/zigzag-conversion/description/
# Date: 2024/10/28
import itertools


def convert(s: str, row_nums: int) -> str:
    """模拟算法
    可以把每个Z看做一个上下翻转的N, 一个字符串的排列可以转变为一个I型与多个倒V型的组合:
    |P |    I|    N|
    |A |  L S|  I G|
    |Y | A  H| R   |
    |P |    I|     |
    如上即为一个I型与两个倒V型的组合, 每个V的容量为 numRows*2 - 2 也就是 (numRows-1) * 2
    设需要倒V的数量为v_nums
    如上面的例子中可以看到:
    当n=14 rn=4时: capacity=6 v_nums=2
    当n=14 rn=5时: capacity=8 v_nums=2 也就是:
    v_nums = (n - numRows + capacity - 1) // capacity
    """
    n = len(s)
    if n < 2 or row_nums < 2:
        return s
    capacity = (row_nums - 1) * 2  # 每个倒V可以容纳的字母数量
    v_nums = (n - row_nums + capacity - 1) // capacity  # 计算倒V的数量, 结果向上取整
    col_nums = 1 + v_nums * (row_nums - 1)  # 矩阵的列数
    matrix = [[0] * col_nums for _ in range(row_nums)]
    x, y = 0, 0  # 定义一个光标用来表示将字符串输入矩阵的位置
    matrix[x][y] = s[0]
    for i in range(1, n):
        if 0 < i % ((row_nums - 1) * 2) < row_nums:
            x, y = x + 1, y  # 光标向下移动
            matrix[x][y] = s[i]
        else:
            x, y = x - 1, y + 1  # 光标向右上方移动
            matrix[x][y] = s[i]
    return "".join(element for row in matrix for element in row if element != 0)


def convert_opt(s: str, row_nums: int) -> str:
    """优化空间的模拟算法
    可以注意到每次在矩阵中的某一行添加字符的时候, 都会添加在上一个字符的右侧, 因此我们可以使用一种压缩矩阵的方法:
    初始化一个空矩阵, 每次向矩阵中添加字符的时候直接添加在该行的末尾即可
    |P    |I    |N
    |A  L |S  I |G
    |Y A  |H R  |
    |P    |I    |
    """
    if row_nums == 1 or row_nums >= len(s):
        return s
    mat = [[] for _ in range(row_nums)]  # 初始化一个空矩阵, 每一行都是一个空列表
    capacity, x = (row_nums - 1) * 2, 0  # capacity每个V可以容纳的字母数量, x表示一个换行的游标
    for i, ch in enumerate(s):
        mat[x].append(ch)
        x += 1 if i % capacity < row_nums - 1 else -1
    # *mat表示解包一个列表, 将mat中的所有子列表连接成一个一维的迭代器（chain object），即展平二维数组。
    # 如果 mat = [['a', 'b'], ['c', 'd']], 则 itertools.chain(*mat) 会生成一个迭代器，按顺序包含 ['a', 'b', 'c', 'd']。
    return ''.join(itertools.chain(*mat))


def convert_opt1(s: str, row_nums: int) -> str:
    """利用索引直接查找
    0     6        12
    1   5 7     11 13
    2 4   8  10    14
    3     9        15
    第一行的索引i的规律为 i%capacity为0
    第二行的索引i的规律为 i%capacity为1或capacity-1
    第三行的索引i的规律为 i%capacity为2或capacity-2
    第三行的索引i的规律为 i%capacity为3或capacity-3
    """
    n, r = len(s), row_nums
    if r == 1 or r >= n:
        return s
    capacity = r * 2 - 2
    ans = []
    for i in range(r):  # 枚举矩阵的行
        for j in range(0, n - i, capacity):  # 枚举每个V型中的起始下标
            ans.append(s[j + i])  # 当前V型中的第一个字符 如上图中的1
            if 0 < i < r - 1 and j + capacity - i < n:
                ans.append(s[j + capacity - i])  # 当前V型中第二个字符 如上图中的5
    return ''.join(ans)


if __name__ == '__main__':
    str = "PAYPALISHIRING"
    print(convert_opt1(str, 4))  # PINALSIGYAHRPI
    print(convert_opt(str, 3))  # PAHNAPLSIIGYIR
    print(convert_opt(str, 5))  # PHASIYIRPLIGAN
